import requests
import json
import sys
from tabulate import tabulate
from colorama import init, Fore, Style
import os
import time
import re
import socket

# 初始化colorama
init()

def find_server_port(start_port=5000, max_port=5100):
    """查找服务端使用的端口"""
    for port in range(start_port, max_port + 1):
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.connect(('localhost', port))
                return port
        except socket.error:
            continue
    return None

def clear_screen():
    """清屏函数"""
    os.system('cls' if os.name == 'nt' else 'clear')

def format_sql(sql):
    """格式化SQL语句显示"""
    # 高亮SQL关键字
    keywords = ['SELECT', 'FROM', 'WHERE', 'AND', 'OR', 'INSERT', 'INTO', 'VALUES', 
                'UPDATE', 'SET', 'DELETE', 'JOIN', 'ON', 'GROUP BY', 'ORDER BY', 
                'HAVING', 'LIMIT', 'OFFSET', 'IN', 'NOT IN', 'EXISTS', 'NOT EXISTS',
                'LIKE', 'NOT LIKE', 'IS NULL', 'IS NOT NULL', 'BETWEEN', 'NOT BETWEEN']
    
    for keyword in keywords:
        sql = re.sub(f'\\b{keyword}\\b', f'{Fore.BLUE}{keyword}{Style.RESET_ALL}', sql, flags=re.IGNORECASE)
    
    # 高亮字符串
    sql = re.sub(r"'([^']*)'", f"{Fore.GREEN}'\\1'{Style.RESET_ALL}", sql)
    sql = re.sub(r'"([^"]*)"', f'{Fore.GREEN}"\\1"{Style.RESET_ALL}', sql)
    
    # 高亮数字
    sql = re.sub(r'\b\d+\b', f'{Fore.YELLOW}\\0{Style.RESET_ALL}', sql)
    
    return sql

def format_table(data):
    """将数据转换为表格格式"""
    if not data:
        return "没有找到数据"
    
    # 如果是字典列表
    if isinstance(data, list) and data and isinstance(data[0], dict):
        headers = data[0].keys()
        # 处理中文字符对齐
        rows = [[str(row[header]) for header in headers] for row in data]
        return tabulate(rows, headers=headers, tablefmt="grid", stralign="center", numalign="center")
    # 如果是单个字典
    elif isinstance(data, dict):
        return tabulate([(k, str(v)) for k, v in data.items()], tablefmt="grid", stralign="center", numalign="center")
    # 其他情况
    else:
        return str(data)

def print_header():
    """打印程序头部信息"""
    clear_screen()
    print(Fore.GREEN + "=" * 60)
    print("欢迎使用MySQL自然语言查询系统！".center(60))
    print("=" * 60 + Style.RESET_ALL)
    print(Fore.YELLOW + "提示：".center(60))
    print("1. 输入 'exit' 或 'quit' 退出程序".center(60))
    print("2. 输入 'clear' 清屏".center(60))
    print("3. 输入 'help' 查看帮助信息".center(60))
    print("=" * 60 + Style.RESET_ALL)

def print_help():
    """打印帮助信息"""
    print(Fore.CYAN + "\n帮助信息：")
    print("1. 查询数据示例：")
    print("   - 查询所有展商信息")
    print("   - 查找来自中国的展商")
    print("   - 统计每个展区的展商数量")
    print("\n2. 添加数据示例：")
    print("   - 添加展区，名称是室内展馆B，位置在主展馆二层")
    print("   - 添加展商，名称是测试公司，国家是中国")
    print("\n3. 修改数据示例：")
    print("   - 将展商ID为1的联系人改为张三")
    print("   - 更新展区ID为1的位置信息")
    print("\n4. 删除数据示例：")
    print("   - 删除ID为1的展商")
    print("   - 删除名称为测试展区的展区" + Style.RESET_ALL)

def main():
    print_header()
    
    # 查找服务端端口
    server_port = find_server_port()
    if not server_port:
        print(Fore.RED + "\n错误：无法找到运行中的服务端，请确保服务端已启动" + Style.RESET_ALL)
        time.sleep(2)
        return
    
    print(Fore.GREEN + f"\n已连接到服务端 (端口: {server_port})" + Style.RESET_ALL)
    
    while True:
        try:
            # 获取用户输入
            query = input("\n" + Fore.GREEN + "请输入您的查询（中文）：" + Style.RESET_ALL).strip()
            
            # 检查特殊命令
            if query.lower() in ['exit', 'quit']:
                print(Fore.GREEN + "\n海内存知己，天涯若比邻，感谢使用，再见！" + Style.RESET_ALL)
                time.sleep(1)
                break
            elif query.lower() == 'clear':
                print_header()
                continue
            elif query.lower() == 'help':
                print_help()
                continue
            
            # 发送请求到服务
            response = requests.post(
                f"http://localhost:{server_port}/query",
                json={"query": query}
            )
            
            # 处理响应
            if response.status_code == 200:
                result = response.json()
                if "queries" in result:
                    for i, query_result in enumerate(result["queries"], 1):
                        print(f"\n" + Fore.CYAN + f"第 {i} 条SQL语句：" + Style.RESET_ALL)
                        print(format_sql(query_result["sql_query"]))
                        print("\n" + Fore.CYAN + "查询结果：" + Style.RESET_ALL)
                        
                        # 格式化输出结果
                        formatted_result = format_table(query_result["result"])
                        print(formatted_result)
                else:
                    print(Fore.RED + "错误：服务器返回的数据格式不正确" + Style.RESET_ALL)
            else:
                print(Fore.RED + f"错误：{response.text}" + Style.RESET_ALL)
                
        except requests.exceptions.ConnectionError:
            print(Fore.RED + "\n错误：无法连接到服务器，请确保服务已启动" + Style.RESET_ALL)
            time.sleep(2)
            break
        except Exception as e:
            print(Fore.RED + f"\n发生错误：{str(e)}" + Style.RESET_ALL)
            time.sleep(2)

if __name__ == "__main__":
    main() 